﻿namespace MetaDrive.VisualStudio.Linq
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using EnvDTE;

    public sealed class CodeElementIterator : IEnumerable<CodeElement>
    {
        private CodeElements codeElements;

        public CodeElementIterator(CodeElements codeElements)
        {
            if (codeElements == null)
                throw new ArgumentNullException("codeElements");

            this.codeElements = codeElements;
        }

        public IEnumerator<CodeElement> GetEnumerator()
        {
            return Enumerate(codeElements).Select(codeElement => codeElement).GetEnumerator();
        }

        IEnumerable<CodeElement> Enumerate(EnvDTE.CodeElements codeElements)
        {
            foreach (CodeElement element in codeElements)
            {
                if (element != null)
                    yield return element;

                CodeElements childrens = null;

                try
                {
                    childrens = element.Children;
                }
                catch (NotImplementedException)
                {
                    childrens = null;
                }

                if (childrens != null)
                {
                    foreach (EnvDTE.CodeElement childElement in Enumerate(childrens))
                    {
                        if (childElement != null)
                            yield return childElement;
                    }
                }
            }
        }

        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }
    }
}
